#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <stdarg.h>
#include <assert.h>
#include <math.h>
#include "ncw.h"

#define INC 10

int verbose = 0;

static void description()
{
    printf("  Description:\n");
    printf("    `nccat' concatenates dimensions <dim> [...] in files <src> [...] and writes\n");
    printf("    result to <dst>. The input files are assumed to have identical structure\n");
    printf("    (apart from the size of the concatenated dimensions). The concatenated\n");
    printf("    dimensions must be the most outward in any variables. The variables that do\n");
    printf("    not depend on the concatenated dimensions are copied from the first source\n");
    printf("    file.\n");
}

static void usage()
{
    printf("  Usage: nccat -d <dim> [...] -i <src> [...] -o <dst> [-v]\n");
    printf("  Options:\n");
    printf("    -v -- verbose\n");
}

static void quit(char* format, ...)
{
    va_list args;

    fflush(stdout);

    fprintf(stderr, "\n\n  ERROR: ");
    va_start(args, format);
    vfprintf(stderr, format, args);
    va_end(args);
    fprintf(stderr, "\n\n");

    exit(1);
}

static void printf_v(char* format, ...)
{
    va_list args;

    if (!verbose)
        return;

    va_start(args, format);
    vprintf(format, args);
    va_end(args);
    fflush(stdout);
}

/** Allocates nj x ni matrix of something. It will be accessed as [i][j].
 * @param ni Outer dimension
 * @param nj Inner dimension
 * @param unitsize Size of one matrix element in bytes
 * @return Matrix
 */
static void* alloc2d(int ni, int nj, size_t unitsize)
{
    size_t size;
    char* p;
    char** pp;
    int i;

    if (ni <= 0 || nj <= 0)
        quit("alloc2d(): invalid size (ni = %d, nj = %d)", ni, nj);

    size = ni * nj;
    if ((p = calloc(size, unitsize)) == NULL) {
        int errno_saved = errno;

        quit("alloc2d(): %s", strerror(errno_saved));
    }
    size = ni * sizeof(void*);
    if ((pp = malloc(size)) == NULL) {
        int errno_saved = errno;

        quit("alloc2d(): %s", strerror(errno_saved));
    }
    for (i = 0; i < ni; i++)
        pp[i] = &p[i * nj * unitsize];

    return pp;
}

static void free2d(void* pp)
{
    void* p;

    p = ((void**) pp)[0];
    free(pp);
    free(p);
}

int main(int argc, char** argv)
{
    int ndims = 0;
    char** dims = NULL;
    int nsrc = 0;
    char** src = NULL;
    char* dst = NULL;
    int i;

    int** dimids = NULL;
    int nvars = 0;
    int* varids = NULL;
    size_t* size = NULL;
    int nvars0;

    if (argc == 1) {
        usage();
        description();
        exit(0);
    }

    i = 1;
    while (i < argc) {
        if (argv[i][0] != '-') {
            usage();
            exit(1);
        }
        if (argv[i][1] == 'd') {
            i++;
            while (i < argc && argv[i][0] != '-') {
                if (ndims % INC == 0)
                    dims = realloc(dims, (ndims + INC) * sizeof(char*));
                dims[ndims] = argv[i];
                ndims++;
                i++;
            }
        } else if (argv[i][1] == 'i') {
            i++;
            while (i < argc && argv[i][0] != '-') {
                if (nsrc % INC == 0)
                    src = realloc(src, (nsrc + INC) * sizeof(char*));
                src[nsrc] = strdup(argv[i]);
                nsrc++;
                i++;
            }
        } else if (argv[i][1] == 'o') {
            i++;
            if (i < argc && argv[i][0] != '-') {
                dst = argv[i];
                i++;
            }
        } else if (argv[i][1] == 'v') {
            verbose = 1;
            i++;
        } else {
            usage();
            exit(1);
        }
    }
    if (ndims == 0 || nsrc < 1 || dst == NULL) {
        usage();
        exit(1);
    }

    printf_v("  src:\n");
    for (i = 0; i < nsrc; ++i)
        printf_v("    \"%s\"\n", src[i]);
    printf_v("  dst = \"%s\"\n", dst);
    printf_v("  merged dimensions:\n");
    for (i = 0; i < ndims; ++i)
        printf_v("    %s\n", dims[i]);

    dimids = alloc2d(nsrc, ndims, sizeof(int));
    for (i = 0; i < nsrc; ++i) {
        int ncid, j;

        ncw_open(src[i], NC_NOWRITE, &ncid);
        for (j = 0; j < ndims; ++j)
            ncw_inq_dimid(src[i], ncid, dims[j], &dimids[i][j]);
        ncw_close(src[i], ncid);
    }

    {
        int ncid;

        ncw_open(src[0], NC_NOWRITE, &ncid);
        ncw_inq_nvars(src[0], ncid, &nvars0);
        varids = malloc(nvars0 * sizeof(int));
        printf_v("  merged variables:\n");
        for (i = 0; i < nvars0; ++i) {
            int dimids_now[NC_MAX_VAR_DIMS];
            char varname[NC_MAX_NAME] = "";
            int j;

            ncw_inq_vardimid(src[0], ncid, i, dimids_now);
            for (j = 0; j < ndims; ++j) {
                if (dimids_now[0] == dimids[0][j]) {
                    varids[nvars] = i;
                    nvars++;
                    ncw_inq_varname(src[0], ncid, i, varname);
                    printf_v("    %s\n", varname);
                    break;
                }
            }
        }
        ncw_close(src[0], ncid);
    }

    printf_v("  creating dst:");
    {
        int ncid_dst, ncid0;
        int ndims0;

        ncw_create(dst, NC_CLOBBER | NC_64BIT_OFFSET, &ncid_dst);

        ncw_open(src[0], NC_NOWRITE, &ncid0);
        ncw_inq_ndims(src[0], ncid0, &ndims0);

        for (i = 0; i < ndims0; ++i) {
            char dimname[NC_MAX_NAME];
            size_t dimlen;
            int dimid;
            int j;

            ncw_inq_dim(src[0], ncid0, i, dimname, &dimlen);
            for (j = 0; j < ndims; ++j) {
                if (i == dimids[0][j]) {
                    int ncid_now, dimid_now;
                    size_t dimlen_now;
                    int s;

                    for (s = 1; s < nsrc; ++s) {
                        ncw_open(src[s], NC_NOWRITE, &ncid_now);
                        ncw_inq_dimid(src[s], ncid_now, dimname, &dimid_now);
                        ncw_inq_dimlen(src[s], ncid_now, dimid_now, &dimlen_now);
                        ncw_close(src[s], ncid_now);
                        dimlen += dimlen_now;
                    }
                    break;
                }
            }
            ncw_def_dim(dst, ncid_dst, dimname, dimlen, &dimid);
        }

        for (i = 0; i < nvars0; ++i) {
            nc_type type;
            int nvardims;
            char varname[NC_MAX_NAME];
            int dimids0[NC_MAX_DIMS], dimids_dst[NC_MAX_DIMS];
            int natts;
            int varid_dst;
            int j;

            ncw_inq_var(src[0], ncid0, i, varname, &type, &nvardims, dimids0, &natts);
            for (j = 0; j < nvardims; ++j) {
                char dimname[NC_MAX_NAME];
                size_t dimlen;

                ncw_inq_dim(src[0], ncid0, dimids0[j], dimname, &dimlen);
                ncw_inq_dimid(dst, ncid_dst, dimname, &dimids_dst[j]);
            }

            ncw_def_var(dst, ncid_dst, varname, type, nvardims, dimids_dst, &varid_dst);
            ncw_copy_atts(src[0], ncid0, i, dst, ncid_dst, varid_dst);
        }

        ncw_close(dst, ncid_dst);
        ncw_close(src[0], ncid0);
    }
    printf_v("\n");

    printf_v("  copying data:");
    size = malloc(nsrc * sizeof(size_t));
    {
        int ncid_dst;

        ncw_open(dst, NC_WRITE, &ncid_dst);

        for (i = 0; i < nvars0; ++i) {
            int j;

            for (j = 0; j < nvars; ++j)
                if (varids[j] == i)
                    break;

            if (j == nvars) {
                int ncid0;

                ncw_open(src[0], NC_NOWRITE, &ncid0);
                ncw_copy_vardata(src[0], ncid0, i, dst, ncid_dst);
                ncw_close(src[0], ncid0);
            } else {
                int ncid;
                char varname[NC_MAX_NAME];
                int s;
                size_t size_total = 0;
                size_t nbytes;
                void* data = NULL;
                void* data_now;
                nc_type type;
                int varid_dst;

                ncw_open(src[0], NC_NOWRITE, &ncid);
                ncw_inq_varname(src[0], ncid, i, varname);
                ncw_close(src[0], ncid);

                size_total = 0;
                for (s = 0; s < nsrc; ++s) {
                    int ncid;
                    int varid;
                    int ndims;
                    int dimids[NC_MAX_DIMS];
                    size_t dimlens[NC_MAX_DIMS];

                    ncw_open(src[s], NC_NOWRITE, &ncid);
                    ncw_inq_varid(src[s], ncid, varname, &varid);
                    ncw_inq_var(src[s], ncid, varid, NULL, &type, &ndims, dimids, NULL);
                    for (j = 0; j < ndims; ++j)
                        ncw_inq_dimlen(src[s], ncid, dimids[j], &dimlens[j]);
                    size[s] = 1;
                    for (j = 0; j < ndims; ++j)
                        size[s] *= dimlens[j];
                    size_total += size[s];
                }

                nbytes = size_total * ncw_sizeof(type);
                if (nbytes > 4294967296)
                    quit("sizeof(%s) = %zu exceeds 4GB", varname, nbytes);
                data = malloc(nbytes);
                if (data == NULL)
                    quit("malloc(): could not allocate memory for variable \"%s\" (size = %zu)", varname, size_total * ncw_sizeof(type));
                data_now = data;
                ncw_inq_varid(dst, ncid_dst, varname, &varid_dst);
                switch (type) {
                case NC_BYTE:
                case NC_CHAR:
                    assert(sizeof(char) == ncw_sizeof(type));
                    for (s = 0; s < nsrc; ++s) {
                        int ncid;
			int varid;

                        ncw_open(src[s], NC_NOWRITE, &ncid);
			ncw_inq_varid(src[s], ncid, varname, &varid);
                        ncw_get_var_text(src[s], ncid, varid, data_now);
                        ncw_close(src[s], ncid);
                        data_now = &((char*) data_now)[size[s]];
                    }
                    ncw_put_var_text(dst, ncid_dst, varid_dst, data);
                    break;
                case NC_SHORT:
                    assert(sizeof(int16_t) == ncw_sizeof(type));
                    for (s = 0; s < nsrc; ++s) {
                        int ncid;
			int varid;

                        ncw_open(src[s], NC_NOWRITE, &ncid);
			ncw_inq_varid(src[s], ncid, varname, &varid);
                        ncw_get_var_short(src[s], ncid, varid, data_now);
                        ncw_close(src[s], ncid);
                        data_now = &((int16_t *) data_now)[size[s]];
                    }
                    ncw_put_var_short(dst, ncid_dst, varid_dst, data);
                    break;
                case NC_INT:
                    assert(sizeof(int32_t) == ncw_sizeof(type));
                    for (s = 0; s < nsrc; ++s) {
                        int ncid;
			int varid;

                        ncw_open(src[s], NC_NOWRITE, &ncid);
			ncw_inq_varid(src[s], ncid, varname, &varid);
                        ncw_get_var_int(src[s], ncid, varid, data_now);
                        ncw_close(src[s], ncid);
                        data_now = &((int32_t *) data_now)[size[s]];
                    }
                    ncw_put_var_int(dst, ncid_dst, varid_dst, data);
                    break;
                case NC_FLOAT:
                    assert(sizeof(float) == ncw_sizeof(type));
                    for (s = 0; s < nsrc; ++s) {
                        int ncid;
			int varid;

                        ncw_open(src[s], NC_NOWRITE, &ncid);
			ncw_inq_varid(src[s], ncid, varname, &varid);
                        ncw_get_var_float(src[s], ncid, varid, data_now);
                        ncw_close(src[s], ncid);
                        data_now = &((float*) data_now)[size[s]];
                    }
                    ncw_put_var_float(dst, ncid_dst, varid_dst, data);
                    break;
                case NC_DOUBLE:
                    assert(sizeof(double) == ncw_sizeof(type));
                    for (s = 0; s < nsrc; ++s) {
                        int ncid;
			int varid;

                        ncw_open(src[s], NC_NOWRITE, &ncid);
			ncw_inq_varid(src[s], ncid, varname, &varid);
                        ncw_get_var_double(src[s], ncid, varid, data_now);
                        ncw_close(src[s], ncid);
                        data_now = &((double*) data_now)[size[s]];
                    }
                    ncw_put_var_double(dst, ncid_dst, varid_dst, data);
                    break;
                }
                free(data);
            }
            printf_v(".");
        }
        ncw_close(dst, ncid_dst);
    }
    printf_v("\n");

    /*
     * clean up
     */
    free(size);
    free(varids);
    free2d(dimids);
    for (i = 0; i < nsrc; ++i)
        free(src[i]);
    free(src);

    return 0;
}
